﻿using System;
using System.Collections.Concurrent;
using System.Text;
using Microsoft.AspNetCore.Http;
using System.Threading.Tasks;
using HzySocket.WebSocket.Interface;
using System.Threading;
using System.Linq;
using System.Net.WebSockets;
using System.IO;

namespace HzySocket.WebSocket.Core
{
    public abstract class HzyWebSocketMiddlewareBase<TSession> : IMiddleware, IWebSocketMiddlewareBase<TSession>
        where TSession : WebSocketSessionBase, new()
    {
        /// <summary>
        /// 不活跃秒数 将被清除连接对象
        /// </summary>
        private static int ActiveTime { get; set; } = 15;

        /// <summary>
        /// 非活跃 服务 Session 清除事件 服务会话主动断开清除事件 传入参数 TAppSession
        /// </summary>
        protected static Action<TSession> InactiveCloseCallBack { get; set; } = null;

        /// <summary>
        /// 连接关闭事件 传入参数 TAppSession
        /// </summary>
        protected static Action<TSession> ConnectCloseCallBack { get; set; } = null;

        /// <summary>
        /// 发送异常事件回调
        /// </summary>
        protected static Action<Exception> SendExceptionCall { get; set; } = null;

        /// <summary>
        /// 新的session连接
        /// </summary>
        protected static Action<TSession> ConnectCall { get; set; } = null;

        /// <summary>
        /// 更新AppSession回调
        /// </summary>
        protected static Action<TSession> UpdateAppSessionCall { get; set; } = null;

        /// <summary>
        /// 错误处理函数
        /// </summary>
        protected static Action<string> ErrorCall { get; set; } = null;

        // 为每个session维护一个消息处理队列和处理任务
        private static readonly ConcurrentDictionary<string, ConcurrentQueue<StringBuilder>> _sessionMessageQueues =
            new ConcurrentDictionary<string, ConcurrentQueue<StringBuilder>>();

        private static readonly ConcurrentDictionary<string, Task> _sessionProcessorTasks =
            new ConcurrentDictionary<string, Task>();

        private static readonly ConcurrentDictionary<string, CancellationTokenSource> _sessionCancellationTokenSources =
            new ConcurrentDictionary<string, CancellationTokenSource>();

        /// <summary>
        /// 消息接收函数
        /// </summary>
        /// <param name="appSession"></param>
        /// <returns></returns>
        public abstract Task ExecuteCommand(TSession appSession);

        public async Task InvokeAsync(HttpContext context, RequestDelegate next)
        {
            //检查是否是 WebSocket 类型请求
            if (!context.WebSockets.IsWebSocketRequest)
            {
                await next.Invoke(context);
                return;
            }

            //链接 地址 后面 必须跟 /ws 路由 或者 /wss
            if (context.Request.IsHttps && !context.Request.Path.ToString().StartsWith("/wss"))
            {
                await next.Invoke(context);
                return;
            }

            if (!context.Request.IsHttps && !context.Request.Path.ToString().StartsWith("/ws"))
            {
                await next.Invoke(context);
                return;
            }

            // 先检查不活跃用户
            await CheckInactiveSessionAsync();

            // 组合会话 id
            var sessionKey = context.Connection.Id + "_" + DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
            var cancellationToken = context.RequestAborted;
            var acceptWebSocketAsync = await context.WebSockets.AcceptWebSocketAsync();

            //检查服务端 WebSocketSession 是否存在
            var appSession = WebSocketSessionContainer<TSession>.GetAppSession(sessionKey) ?? new TSession();
            var isFirstConnection = string.IsNullOrWhiteSpace(appSession.SessionKey);
            appSession.HttpContext = context;
            appSession.WebSocket = acceptWebSocketAsync;
            appSession.LastActiveTime = DateTime.Now;
            appSession.GetCancellationToken = cancellationToken;
            appSession.SessionKey = sessionKey;
            if (isFirstConnection)
            {
                appSession.StartTime = DateTime.Now;
                appSession.SendExceptionCall = SendExceptionCall;
                ConnectCall?.Invoke(appSession);
            }
            else
            {
                UpdateAppSessionCall?.Invoke(appSession);
            }

            _ = Task.Run(async () =>
            {
                if (isFirstConnection)
                {
                    await appSession.SendAsync(new { status = 1, msg = $"WS {sessionKey} 请注册用户!" });
                }
                else
                {
                    await appSession.SendAsync(new { status = 1, msg = $"WS {sessionKey} 用户已存在!" });
                }
            }, cancellationToken);

            WebSocketSessionContainer<TSession>.AddOrUpdateAppSession(appSession);

            // 初始化队列
            InitializeSessionQueue(appSession);

            try
            {
                //检查和处理客户端发送的数据
                while (!acceptWebSocketAsync.CloseStatus.HasValue &&
                       (acceptWebSocketAsync.State == WebSocketState.Open ||
                        acceptWebSocketAsync.State == WebSocketState.CloseReceived))
                {
                    if (cancellationToken.IsCancellationRequested)
                    {
                        Console.WriteLine("Cancellation requested.");
                        break;
                    }

                    // 记录活跃时间
                    appSession.LastActiveTime = DateTime.Now;

                    // 接收数据
                    var data = await ReceiveStringAsync(acceptWebSocketAsync, sessionKey, cancellationToken);
                    if (data == null || data.Length == 0) continue;

                    // 将消息添加到对应session的队列中
                    EnqueueMessage(sessionKey, data);
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex);
                ErrorCall?.Invoke($"Error processing WebSocket message: {ex}");
            }
            finally
            {
                await CloseSocketAsync(acceptWebSocketAsync, sessionKey);

                try
                {
                    //连接断开后移除 session 对象
                    ConnectCloseCallBack?.Invoke(appSession);
                }
                catch (Exception e)
                {
                    Console.WriteLine(e);
                }

                // 清理该session的资源
                CleanupSessionResources(sessionKey);
                WebSocketSessionContainer<TSession>.RemoveSession(w => w.SessionKey == sessionKey);
                // 释放掉这个 socket
                acceptWebSocketAsync?.Dispose();
            }
            //参考了：https://www.cnblogs.com/besuccess/p/7043885.html     
        }

        private static async Task CloseSocketAsync(System.Net.WebSockets.WebSocket acceptWebSocketAsync,
            string sessionKey)
        {
            // 确保WebSocket连接被关闭
            try
            {
                await acceptWebSocketAsync.CloseAsync(WebSocketCloseStatus.NormalClosure, "Closing",
                    CancellationToken.None);
            }
            catch (Exception ex)
            {
                ErrorCall?.Invoke($"Error closing WebSocket: {ex}");
                acceptWebSocketAsync.Abort();
            }
            finally
            {
                WebSocketSessionContainer<TSession>.RemoveSession(w => w.SessionKey == sessionKey);
            }
        }

        private static async Task<StringBuilder> ReceiveStringAsync(System.Net.WebSockets.WebSocket socket,
            string sessionKey,
            CancellationToken ct = default(CancellationToken))
        {
            // 检查WebSocket状态
            if (socket.State != WebSocketState.Open && socket.State != WebSocketState.CloseReceived)
            {
                return null;
            }

            // 使用ArrayPool来减少内存分配
            var bufferArray = System.Buffers.ArrayPool<byte>.Shared.Rent(1024 * 4);
            var buffer = new ArraySegment<byte>(bufferArray);

            try
            {
                using (var ms = new MemoryStream())
                {
                    try
                    {
                        WebSocketReceiveResult webSocketReceiveResult;
                        do
                        {
                            // 检查取消令牌
                            if (ct.IsCancellationRequested)
                            {
                                throw new OperationCanceledException();
                            }

                            // 检查WebSocket状态
                            if (socket.State != WebSocketState.Open && socket.State != WebSocketState.CloseReceived)
                            {
                                return null;
                            }

                            ct.ThrowIfCancellationRequested();

                            // 添加接收超时控制，避免长时间阻塞
                            using (var cts = CancellationTokenSource.CreateLinkedTokenSource(ct))
                            {
                                cts.CancelAfter(TimeSpan.FromSeconds(30)); // 30秒超时

                                webSocketReceiveResult = await socket.ReceiveAsync(buffer, cts.Token);

                                // 检查消息大小，防止内存耗尽攻击
                                if (ms.Length + webSocketReceiveResult.Count > 1024 * 1024) // 限制最大1MB
                                {
                                    throw new InvalidOperationException("Message too large");
                                }

                                ms.Write(buffer.Array, buffer.Offset, webSocketReceiveResult.Count);
                            }
                        } while (!webSocketReceiveResult.EndOfMessage);

                        // 检查消息类型
                        if (webSocketReceiveResult.MessageType != WebSocketMessageType.Text)
                        {
                            return null;
                        }

                        // 重置流位置并读取内容
                        ms.Seek(0, SeekOrigin.Begin);
                        // 直接从内存流中读取数据并转换为字符串，然后创建StringBuilder
                        var bytes = ms.ToArray();
                        var text = Encoding.UTF8.GetString(bytes);
                        return new StringBuilder(text);
                    }
                    catch (Exception ex)
                    {
                        // 其他异常记录日志
                        ErrorCall?.Invoke($"Error receiving message: {ex}");
                        Console.WriteLine(ex);
                        throw;
                    }
                }
            }
            finally
            {
                // 归还缓冲区到池中
                System.Buffers.ArrayPool<byte>.Shared.Return(bufferArray);
            }
        }

        /// <summary>
        /// 检查不活跃Session信息
        /// </summary>
        /// <returns></returns>
        protected virtual async Task CheckInactiveSessionAsync()
        {
            //移除不活跃的SocketSession对象
            var tAppSessionList = WebSocketSessionContainer<TSession>.GetAppSessions(w =>
                (DateTime.Now - w.LastActiveTime).TotalSeconds >= ActiveTime);
            var webSocketSessionBases = tAppSessionList as TSession[] ?? tAppSessionList.ToArray();
            if (!webSocketSessionBases.Any()) return;

            try
            {
                foreach (var item in webSocketSessionBases)
                {
                    InactiveCloseCallBack?.Invoke(item);

                    // 主动关闭WebSocket连接以释放资源
                    try
                    {
                        await item.WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Closing",
                            CancellationToken.None);
                    }
                    catch (Exception ex)
                    {
                        item.WebSocket?.Abort();
                        ErrorCall?.Invoke($"Error closing WebSocket for inactive session: {ex}");
                    }

                    // 清理该session的资源
                    CleanupSessionResources(item.SessionKey);
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex);
                ErrorCall?.Invoke($"CheckInactiveSessionAsync error: {ex}");
            }

            WebSocketSessionContainer<TSession>.RemoveSession(w =>
                (DateTime.Now - w.LastActiveTime).TotalSeconds >= ActiveTime);
        }

        #region 消息队列

        private void InitializeSessionQueue(TSession appSession)
        {
            var messageQueue = new ConcurrentQueue<StringBuilder>();
            var cts = new CancellationTokenSource();

            _sessionMessageQueues.TryAdd(appSession.SessionKey, messageQueue);
            _sessionCancellationTokenSources.TryAdd(appSession.SessionKey, cts);

            // 启动该session的消息处理任务
            var processorTask = Task.Run(async () => await ProcessSessionMessages(appSession, cts.Token), cts.Token);
            _sessionProcessorTasks.TryAdd(appSession.SessionKey, processorTask);
        }

        private void EnqueueMessage(string sessionKey, StringBuilder data)
        {
            if (_sessionMessageQueues.TryGetValue(sessionKey, out var messageQueue))
            {
                messageQueue.Enqueue(data);
            }
        }

        private async Task ProcessSessionMessages(TSession appSession,
            CancellationToken cancellationToken)
        {
            var sessionKey = appSession.SessionKey;
            // 添加最大队列长度监控，防止内存溢出
            const int maxQueueLength = 1000000;

            while (!cancellationToken.IsCancellationRequested)
            {
                try
                {
                    appSession.LastActiveTime = DateTime.Now;
                    if (_sessionMessageQueues.TryGetValue(sessionKey, out var messageQueue))
                    {
                        // 检查队列长度，防止消息堆积过多
                        if (messageQueue.Count > maxQueueLength)
                        {
                            ErrorCall?.Invoke(
                                $" 【检查队列长度，防止消息堆积过多】 Message queue for session {sessionKey} is too long: {messageQueue.Count}");
                            // 可以选择丢弃旧消息或采取其他策略
                        }

                        if (messageQueue.TryDequeue(out var data))
                        {
                            if (RequiresOrderedProcessing(data))
                            {
                                // 顺序处理消息
                                await ProcessMessage(appSession, data);
                            }
                            else
                            {
                                // 并行处理不需要顺序的消息
                                _ = Task.Run(async () => await ProcessMessage(appSession, data), cancellationToken);
                            }
                        }
                        else
                        {
                            // 没有消息时短暂休眠，避免CPU空转 // 使用更智能的等待机制
                            await Task.Delay(Math.Min(100, messageQueue.Count > 0 ? 50 : 100), CancellationToken.None);
                        }
                    }
                    else
                    {
                        // 队列已被移除，退出循环
                        break;
                    }
                }
                catch (Exception ex)
                {
                    Console.WriteLine(ex);
                    ErrorCall?.Invoke($"ProcessSessionMessages error: {ex}");
                    break;
                }
            }
        }

        private async Task ProcessMessage(TSession appSession, StringBuilder data)
        {
            appSession.Data = data;
            appSession.LastActiveTime = DateTime.Now;

            try
            {
                await this.ExecuteCommand(appSession);
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex);
                // 记录异常，避免Task.Run异常丢失
                ErrorCall?.Invoke($"ExecuteCommand error: {ex}");
            }
            finally
            {
                appSession.LastActiveTime = DateTime.Now;
                WebSocketSessionContainer<TSession>.AddOrUpdateAppSession(appSession);
            }
        }

        private static void CleanupSessionResources(string sessionKey)
        {
            // 取消该session的消息处理任务
            if (_sessionCancellationTokenSources.TryRemove(sessionKey, out var cts))
            {
                cts.Cancel();
                cts.Dispose();
            }

            // 移除队列和任务引用
            _sessionMessageQueues.TryRemove(sessionKey, out _);
            _sessionProcessorTasks.TryRemove(sessionKey, out _);
        }

        #endregion

        /// <summary>
        /// 判断消息是否需要顺序处理的辅助方法
        /// </summary>
        /// <param name="data"></param>
        /// <returns></returns>
        protected virtual bool RequiresOrderedProcessing(StringBuilder data)
        {
            // 默认不按顺序执行
            return data != null && data.Length > 0;
        }

        /// <summary>
        /// 设置会话活动时间
        /// </summary>
        /// <param name="sec"></param>
        public void SetActiveTime(int sec)
        {
            ActiveTime = sec;
        }

        /// <summary>
        /// 获取会话活动时间
        /// </summary>
        /// <returns></returns>
        public int GetActiveTime()
        {
            return ActiveTime;
        }
    }
}